# *_*coding:utf-8 *_*
# @Author : Reggie
# @Time : 2022/1/6 15:36
import json
import logging
import os
import shlex
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import paramiko

ENVIRON = {
    "OFMT": "JSON",
    "LANG": "en_US.UTF-8;"
}
LOG_LEVEL = "INFO"

env_debug = True if os.environ.get("DEBUG") else False
if env_debug:
    LOG_LEVEL = "DEBUG"

nolog = True if os.environ.get("nolog") else False
if not nolog:
    root_logger = logging.getLogger()
    stream_handle = logging.StreamHandler()
    formatter = logging.Formatter("%(asctime)s -- %(levelname)s --"
                                  " %(message)s (%(funcName)s in %(filename)s %(lineno)d)")
    stream_handle.setFormatter(formatter)
    root_logger.addHandler(stream_handle)
    root_logger.setLevel(LOG_LEVEL)


def run_command(cmd,
                shell=False,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                stdin=subprocess.PIPE,
                input=None,
                cmd_env=None,
                raise_error=False,
                json_out=False,
                encode="utf-8"):
    out, err, rc = {}, "", 1

    if shell:
        if isinstance(cmd, (list, tuple)):
            cmd = " ".join([str(i) for i in cmd])
    else:
        if isinstance(cmd, str):
            cmd = shlex.split(cmd)
        elif isinstance(cmd, (list, tuple)):
            cmd = [str(i) for i in cmd]
    cmd_str = " ".join(cmd) if isinstance(cmd, (list, tuple)) else cmd

    env = dict(os.environ)
    if cmd_env:
        for k in cmd_env.keys():
            env[k] = cmd_env[k]

    try:
        logging.debug(f"Running command: {cmd_str}")
        p = subprocess.Popen(args=cmd, shell=shell, stdout=stdout, stderr=stderr, stdin=stdin, env=env)
        out, err = p.communicate(input=input)
        rc = p.returncode
    except Exception as e:
        err = f"Run Command Error: {e}"
        if raise_error:
            raise Exception(e)
        return out, err, rc

    if encode:
        out = out.decode(encode)
        err = err.decode(encode)

    if json_out:
        try:
            out = json.loads(out)
        except Exception as e:
            err = f"JSON Loads Error: {e}"
            if raise_error:
                raise Exception(e)
            return out, err, rc

    return out, err, rc


class SSHConnection(object):
    def __init__(self, host, port=22, username='root', pwd='user@dev', use_callback=False):
        self.host = host
        self.port = port
        self.username = username
        self.pwd = pwd
        self.use_callback = use_callback
        self.__transport = None
        self._ssh = None
        self._sftp = None

    def __enter__(self):
        if self.__transport is not None:
            raise RuntimeError("Already connected")
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def connect(self):
        if self.__transport is not None and self.__transport.active:
            return
        transport = paramiko.Transport((self.host, self.port))
        transport.connect(username=self.username, password=self.pwd)
        self.__transport = transport

    def close(self):
        self.__transport.close()
        self.__transport = None

    def translate_byte(self, B):
        B = float(B)
        KB = float(1024)
        MB = float(KB ** 2)
        GB = float(MB ** 2)
        TB = float(GB ** 2)
        if B < KB:
            return '{} {}'.format(B, 'bytes' if B > 1 else 'byte')
        elif KB < B < MB:
            return '{:.2f} KB'.format(B / KB)
        elif MB < B < GB:
            return '{:.2f} MB'.format(B / MB)
        elif GB < B < TB:
            return '{:.2f} GB'.format(B / GB)
        else:
            return '{:.2f} TB'.format(B / TB)

    def call_back(self, transferred, toBeTransferred, suffix=''):
        bar_len = 100
        filled_len = int(round(bar_len * transferred / float(toBeTransferred)))
        percents = round(100.0 * transferred / float(toBeTransferred), 1)
        bar = '\033[32;1m%s\033[0m' % '=' * filled_len + '-' * (bar_len - filled_len)
        sys.stdout.write('[%s] %s%s %s\r\n' % (bar, '\033[32;1m%s\033[0m' % percents, '%', suffix))
        sys.stdout.flush()

    def sftp_down_file(self, server_path, local_path):
        if self._sftp:
            sftp = self._sftp
        else:
            sftp = paramiko.SFTPClient.from_transport(self.__transport)
            self._sftp = sftp
        result = sftp.get(server_path, local_path, callback=self.call_back if self.use_callback else None)
        return result

    def sftp_upload_file(self, local_path, server_path):
        if self._sftp:
            sftp = self._sftp
        else:
            sftp = paramiko.SFTPClient.from_transport(self.__transport)
            print("init")
            self._sftp = sftp
        result = sftp.put(local_path, server_path, callback=self.call_back if self.use_callback else None)
        return result

    def cmd(self,
            command,
            encoding="utf-8",
            bufsize=-1,
            timeout=None,
            get_pty=False,
            environment=None,
            ):
        if self._ssh:
            ssh = self._ssh
        else:
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh._transport = self.__transport
            self._ssh = ssh
        # 执行命令
        stdin, stdout, stderr = ssh.exec_command(
            command,
            bufsize=bufsize,
            timeout=timeout,
            get_pty=get_pty,
            environment=environment,
        )
        # 获取命令结果
        stdout = stdout.read().decode(encoding)
        stderr = stderr.read().decode(encoding)
        return stdout, stdin, stderr


def single_install(ip):
    start = time.time()
    logging.info(f"start {ip}")
    UxsName = "uxs-mgr-1.2.4"

    install_rpm_command = f"""OldPath=`pwd`
    UxsName="{UxsName}"
    echo "UxsName: ${{UxsName}}"
    echo "OldPath: ${{UxsName}}"
    cd /root
    rm -rf ${{UxsName}}
    tar zxvf ${{UxsName}}.tar.gz
    cd ${{UxsName}}
    rpm -iUh ./*.rpm
    #yum list | grep uxs
    rpm -qa | grep uxs
    cd /root
    cd $OldPath
    """
    restart_service_command = """systemctl daemon-reload"""

    with SSHConnection(host=ip) as client:
        logging.info(f"Step 1 install rpm")
        stdout, stdin, stderr = client.cmd(install_rpm_command)
        print(stdout)
        print(stderr)
        logging.info(f"Step 2 restart service")
        stdout, stdin, stderr = client.cmd(restart_service_command)
        print(stdout)
        print(stderr)

    def rcmd(cmd):
        cmd = cmd.strip()
        with SSHConnection(host=ip) as client:
            logging.info(f"Step 2 restart service")
            stdout, stdin, stderr = client.cmd(cmd)
            print(stdout)
            print(stderr)
        return cmd

    restart_service_command = """systemctl restart uxs-monitor.service
    systemctl restart uxs-mounter.service
    systemctl restart uxs-watchdog.service
    systemctl restart uxs-agent.service
    systemctl restart uxs-dashboard.service"""
    service_command_set = restart_service_command.split("\n")

    executor = ThreadPoolExecutor(max_workers=len(service_command_set))
    tasks = [executor.submit(rcmd, cmd) for cmd in service_command_set]

    for future in as_completed(tasks):
        data = future.result()
        print(f"{data}")
    return f"{ip} finished. Cost: {time.time() - start}"


def main(ips):
    """
    :param ips: str     "172.16.120.134 172.16.120.135 172.16.120.136"
                    list    ["172.16.120.134", "172.16.120.135", "172.16.120.136"]
    :return:
    """
    if isinstance(ips, str):
        ip_list = ips.split()
    elif isinstance(ips, list):
        ip_list = ips
    else:
        raise ValueError(f"{ips} type error. Must be list or str")
    if not ips:
        raise Exception(f"{ips} can not be empty")
    print(ip_list)

    # for ip in ip_list:
    #     single_install(ip)

    executor = ThreadPoolExecutor(max_workers=10)
    tasks = [executor.submit(single_install, ip) for ip in ip_list]

    for future in as_completed(tasks):
        data = future.result()
        print(f"{data}")


if __name__ == '__main__':
    ips = "10.200.102.14 10.200.102.19 10.200.102.25"
    # ips = "10.200.102.11 10.200.102.12 10.200.102.15"
    # ips = "172.16.120.134 172.16.120.135 172.16.120.136"
    # ips = "10.200.104.97 10.200.104.98 10.200.104.99"
    # ips = "10.200.2.101 10.200.2.60 10.200.2.100"
    # ips = "172.16.130.171 172.16.130.184 172.16.130.188"
    # ips = "10.200.2.30 10.200.2.31 10.200.2.32"
    # ips = "10.200.2.60 10.200.2.100 10.200.2.101"
    # ips = "10.200.104.91 10.200.104.92 10.200.104.93 10.200.104.95 10.200.104.96"
    main(ips)
